{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "knowledge-neurons.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xg1Zpk9rntkx"
      },
      "source": [
        "!pip install knowledge-neurons"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3yyA9QwSrnf6"
      },
      "source": [
        "# make sure you have a gpu available\n",
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x0Ll8Hi-n2qq"
      },
      "source": [
        "from knowledge_neurons import (\n",
        "    KnowledgeNeurons,\n",
        "    initialize_model_and_tokenizer,\n",
        "    model_type,\n",
        ")\n",
        "import random\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "# setup model, tokenizer + kn class\n",
        "MODEL_NAME = \"bert-base-multilingual-uncased\"\n",
        "model, tokenizer = initialize_model_and_tokenizer(MODEL_NAME)\n",
        "kn = KnowledgeNeurons(model, tokenizer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3T43iPMDrCpQ"
      },
      "source": [
        "TEXT = \"Sarah was visiting [MASK], the capital of france\"\n",
        "GROUND_TRUTH = \"paris\"\n",
        "BATCH_SIZE = 10\n",
        "STEPS = 20\n",
        "\n",
        "ENG_TEXTS = [\n",
        "    \"Sarah was visiting [MASK], the capital of france\",\n",
        "    \"The capital of france is [MASK]\",\n",
        "    \"[MASK] is the capital of france\",\n",
        "    \"France's capital [MASK] is a hotspot for romantic vacations\",\n",
        "    \"The eiffel tower is situated in [MASK]\",\n",
        "    \"[MASK] is the most populous city in france\",\n",
        "    \"[MASK], france's capital, is one of the most popular tourist destinations in the world\",\n",
        "]\n",
        "FRENCH_TEXTS = [\n",
        "    \"Sarah visitait [MASK], la capitale de la france\",\n",
        "    \"La capitale de la france est [MASK]\",\n",
        "    \"[MASK] est la capitale de la france\",\n",
        "    \"La capitale de la France [MASK] est un haut lieu des vacances romantiques\",\n",
        "    \"La tour eiffel est située à [MASK]\",\n",
        "    \"[MASK] est la ville la plus peuplée de france\",\n",
        "    \"[MASK], la capitale de la france, est l'une des destinations touristiques les plus prisées au monde\",\n",
        "]\n",
        "TEXTS = ENG_TEXTS + FRENCH_TEXTS\n",
        "\n",
        "P = 0.5 # sharing percentage"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ie4f0la4n9kn"
      },
      "source": [
        "# find the neurons separately for each language, then all the languages together\n",
        "\n",
        "refined_neurons_eng = kn.get_refined_neurons(\n",
        "    ENG_TEXTS,\n",
        "    GROUND_TRUTH,\n",
        "    p=P,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    steps=STEPS,\n",
        ")\n",
        "refined_neurons_fr = kn.get_refined_neurons(\n",
        "    FRENCH_TEXTS,\n",
        "    GROUND_TRUTH,\n",
        "    p=P,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    steps=STEPS,\n",
        ")\n",
        "refined_neurons = kn.get_refined_neurons(\n",
        "    TEXTS,\n",
        "    GROUND_TRUTH,\n",
        "    p=P,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    steps=STEPS,\n",
        ")\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V7oSbJUssDas"
      },
      "source": [
        "# layer_no, neuron_idx\n",
        "refined_neurons"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_oenOqTarWh8"
      },
      "source": [
        "# how many neurons are shared between the french prompts and the english ones?\n",
        "print(\"N french neurons: \", len(refined_neurons_fr))\n",
        "print(\"N english neurons: \", len(refined_neurons_eng))\n",
        "shared_neurons = [i for i in refined_neurons_eng if i in refined_neurons_fr]\n",
        "print(f\"N shared neurons: \", len(shared_neurons))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z-6-RO3erVPG"
      },
      "source": [
        "# suppress the refined neurons that we found for a new piece of text expressing the fact\n",
        "TEXT = \"The louvre art museum is located in the city of [MASK]\"\n",
        "GROUND_TRUTH = \"paris\"\n",
        "\n",
        "print(\"\\nSuppressing refined neurons: \\n\")\n",
        "results_dict, unpatch_fn = kn.suppress_knowledge(\n",
        "    TEXT, GROUND_TRUTH, refined_neurons\n",
        ")\n",
        "\n",
        "# suppressing the same amount of random neurons has a much smaller effect\n",
        "print(\"\\nSuppressing random neurons: \\n\")\n",
        "random_neurons = [\n",
        "    [\n",
        "        random.randint(0, model.config.num_hidden_layers - 1),\n",
        "        random.randint(0, model.config.intermediate_size - 1),\n",
        "    ]\n",
        "    for i in range(len(refined_neurons))\n",
        "]\n",
        "results_dict, unpatch_fn = kn.suppress_knowledge(\n",
        "    TEXT, GROUND_TRUTH, random_neurons\n",
        ")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JZ2Pb02Rs8JA"
      },
      "source": [
        "# can we use the neurons found *just* with the french prompts to suppress the\n",
        "# correct english fact being produced?\n",
        "print(\n",
        "    \"\\nSuppressing refined neurons (found by french text) using english prompt: \\n\"\n",
        ")\n",
        "results_dict, unpatch_fn = kn.suppress_knowledge(\n",
        "    TEXT, GROUND_TRUTH, refined_neurons_fr\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XNnrNbaWtIt-"
      },
      "source": [
        "# we can also \"enhance\" a fact\n",
        "print(\"\\nEnhancing refined neurons: \\n\")\n",
        "results_dict, unpatch_fn = kn.enhance_knowledge(\n",
        "    TEXT, GROUND_TRUTH, refined_neurons\n",
        ")\n",
        "\n",
        "print(\"\\nEnhancing random neurons: \\n\")\n",
        "results_dict, unpatch_fn = kn.enhance_knowledge(\n",
        "    TEXT, GROUND_TRUTH, random_neurons\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zUaxoApktg9V"
      },
      "source": [
        "# or we can make the model think the capital of france is London!\n",
        "\n",
        "# edit the weights of the output ff layer at the refined neurons (replacing them with the word embedding of 'target') + test the effect\n",
        "# setting undo_modification to false leaves the modification in - you can play with the model, and then use 'unpatch_fn' to undo it later.\n",
        "results_dict, unpatch_fn = kn.edit_knowledge(\n",
        "    TEXT, target=\"london\", neurons=refined_neurons, undo_modification=False,\n",
        ")\n",
        "edited_model = kn.model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a6en_KmIunX8"
      },
      "source": [
        "encoded_input = kn.tokenizer(\"The louvre art museum is located in the city of [MASK]\", return_tensors=\"pt\").to(kn.device)\n",
        "mask_idx = torch.where(\n",
        "                encoded_input[\"input_ids\"][0] == kn.tokenizer.mask_token_id\n",
        "            )[0].item()\n",
        "outputs = edited_model(**encoded_input)\n",
        "        \n",
        "probs = F.softmax(outputs.logits[:, mask_idx, :], dim=-1)\n",
        "_, argmax_id = [i.item() for i in probs.max(dim=-1)]\n",
        "argmax_str = kn.tokenizer.decode([argmax_id])\n",
        "print(f'The louvre art museum is located in the city of {argmax_str}')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}